from fabric import Connection
from concurrent.futures import ThreadPoolExecutor, as_completed
import argparse

def parse():
    parser = argparse.ArgumentParser(description='Distributed running')
    parser.add_argument('--kill', type=int, default=0, help='Set to 1 to kill processes.')
    args = parser.parse_args()
    return args

args = parse()

# 节点信息，包括远程主机的IP地址、用户名、密码和用户的NetSenseML目录
nodes = {
    "192.168.1.154": {"user": "d", "password": "d", "remote_directory": "/home/d/SparseML"},
    "192.168.1.169": {"user": "dd", "password": "dd", "remote_directory": "/home/dd/sparseML"},
    "192.168.1.157": {"user": "ddd", "password": "ddd", "remote_directory": "/home/ddd/SparseML"},
    "192.168.1.108": {"user": "dddd", "password": "dddd", "remote_directory": "/home/dddd/SparseML"},
    # "192.168.1.107": {"user": "ddddd", "password": "ddddd", "remote_directory": "/home/ddddd/SparseML"},
    # "192.168.1.232": {"user": "dddddd", "password": "dddddd", "remote_directory": "/home/dddddd/SparseML"},
    # "192.168.1.199": {"user": "ddddddd", "password": "ddddddd", "remote_directory": "/home/ddddddd/SparseML"},
    # "192.168.1.248": {"user": "dddddddd", "password": "dddddddd", "remote_directory": "/home/dddddddd/SparseML"}
}

# 当前主机的 rank 和 world_size
world_size = 4

# 定义剪枝比率从 0.9 到 0，步长为 0.1
pruning_amounts = [round(x, 1) for x in list(reversed([i * 0.1 for i in range(9)]))]
pruning_amounts.append(0.99)
# pruning_amounts = [0, 0.1]


# 定义命令模板
command_template = (
    "cd {remote_directory} && make dist-run-resnet world_size={world_size} rank={rank} "
    "model_name=vit-base16 dataset=cifar100 pruning_amount={pruning_amount} hook=allreduce "
    "compression_ratio=0 threshold=0"
)

# 如果是kill模式，覆盖命令模板
if args.kill == 1:
    command_template = 'cd {remote_directory} && ./kill_port.sh 8003'

# 执行 SSH 命令的函数
def fabric_execute_command(hostname, username, password, remote_directory, rank):
    try:
        # 创建 Fabric 连接
        conn = Connection(host=hostname, user=username, connect_kwargs={"password": password})
        
        # 确定是否隐藏输出
        hide = True if username != 'd' else False
        
        # 针对每个剪枝比率执行命令
        for pruning_amount in pruning_amounts:
            command = command_template.format(
                remote_directory=remote_directory,
                world_size=world_size,
                rank=rank,
                pruning_amount=pruning_amount
            )
            print(f"Executing on {hostname}: {command}")
            result = conn.run(command, hide=hide, pty=True)
            print(f"[{hostname}] {result.stdout.strip()}")
            if result.stderr:
                print(f"[{hostname} ERROR] {result.stderr.strip()}")

    except Exception as e:
        print(f"Error executing commands on {hostname}: {e}")
    finally:
        conn.close()

# 并行执行远程和本地命令
def run_commands_in_parallel():
    with ThreadPoolExecutor() as executor:
        futures = []
        rank = 0
        for hostname, credentials in nodes.items():
            username = credentials['user']
            password = credentials['password']
            remote_directory = credentials['remote_directory']

            # 提交远程任务
            futures.append(executor.submit(fabric_execute_command, hostname, username, password, remote_directory, rank))
            rank += 1

        # 等待所有任务完成
        for future in as_completed(futures):
            try:
                future.result()
            except Exception as e:
                print(f"Task failed: {e}")

# 执行
run_commands_in_parallel()
print("Distributed computation commands executed successfully.")